#!/usr/bin/env python
# -*- coding: utf-8 -*-
from langchain.prompts import ChatPromptTemplate

# 创建关键提示词
def create_sql_template():
    return ChatPromptTemplate.from_template(
        "你是一个只负责生成 SQLite 语句的助手。你必须严格按照用户的需求，仅输出 SQLite 语句本身，不包含任何解释、注释或多余内容。\n"
        "所有输出都必须是完整可运行SQLite 语句\n"
        "问题：{q}\n"
        "数据库表结构：{table_define}\n"
        "注：1.select 出来的字段不要增加别称。最小化输出列，不要添加无关语义的列。\n"
        # "2.如果生成的SQL条件部分是年度busi_time=yyyy-MM或者月度busi_time=yyyy格式，需要进行汇总处理。\n"
        "3.busi_time如果是年度和月度统计使用strftime函数，\n"
        " 如：月度统计strftime('%Y-%m', busi_time)='2024-01'\n"
        "   年度统计strftime('%Y', busi_time)='2024'\n"
        " 处理当天数据使用date函数，如：date(busi_time) = '2024-01-01' \n"
        # "4.数据库表结构，除busi_time外，如果其他2个以上(包含）的string字段 \n"
        # " 例如，库表结构包括sub_name string,top_name string,top_cnt int三个字段，如果生成的SQL，只有sub_name和top_cnt \n"
        # " 那么应该按照如下范式：select sub_name,sum(top_cnt) from abc group by sub_name \n"
        # " 增加group by进行分类汇总，如果涉及平均值计算对应字段增加avg() ，如果涉及累加计算对应字段增加sum()\n"
        # " 正确样式：SELECT colume_1,sum(colume_2) from table_1 where colume_3='xyz' GROUP BY colume_1 \n"
        # " 错误样式：SELECT colume_1,colume_2 from table_1 where colume_3='xyz' GROUP BY colume_1 \n"
        "{extra_hint}"
        "{example_hint}"
    )

def debug_print(debug, *args, **kwargs):
    """仅在 DEBUG 模式下打印"""
    if debug :
        print(*args, **kwargs)

# sql后处理
def process_sql_answer(answer):
    try:
        ind = answer.index("SELECT")
        answer = answer[ind:]
        ind = answer.index(';')
        answer = answer[:ind + 1]
    except ValueError:
        pass

    # 去除引号
    if answer[0] in ['"', "'"]:
        answer = answer[1:]
    if answer[-1] == '"':
        answer = answer[:-1]
    if answer[-1] != ';':
        answer += ';'
    if answer.startswith("```"):
        answer = answer.replace("```", "")
    
    time.sleep(3)

    return answer
